from datetime import datetime, timedelta

from PyQt5.QtCore import QAbstractItemModel, QModelIndex
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QIcon

from rss_reader.database.models import Feed, Article
from rss_reader.logger import logger

from PyQt5.QtCore import QAbstractItemModel, Qt, QModelIndex
from datetime import datetime, timedelta


class FeedTreeModel(QAbstractItemModel):
    def __init__(self, user, session_factory, parent=None):
        logger.debug("init")
        super().__init__(parent)
        self.user = user
        self.session_factory = session_factory
        self._setup_model_data()

    def _setup_model_data(self):
        """初始化树形结构数据"""
        logger.debug("start")
        with self.session_factory() as session:
            # 获取所有Feed并按分类分组
            all_feeds = session.query(Feed).order_by(Feed.title).all()

            # 组织分类数据 {分类名称: [feed1, feed2]}
            self.category_feeds = {}
            self.uncategorized_feeds = []

            for feed in all_feeds:
                if feed.category and feed.category.strip():
                    self.category_feeds.setdefault(feed.category, []).append(feed)
                else:
                    self.uncategorized_feeds.append(feed)
            logger.debug(self.category_feeds)
            logger.debug(self.uncategorized_feeds)

            # 计算未读数量
            self.unread_counts = {
                feed.id: session.query(Article).filter(
                    Article.feed_id == feed.id, Article.read == False
                ).count()
                for feed in all_feeds
            }

            # 计算今日文章数
            today = datetime.now() - timedelta(days=1)
            self.today_count = session.query(Article)\
                .filter(Article.published >= today)\
                .count()

        # 计算收藏文章数
        self.starred_count = session.query(Article)\
            .filter(Article.starred == True)\
            .count()


    def columnCount(self, parent=QModelIndex()):
        return 1


    def rowCount(self, parent=QModelIndex()):
        if not parent.isValid():
            # 根节点：3个固定项 + 分类数量 + (未分类项如果存在)
            return 3 + len(self.category_feeds) + (1 if self.uncategorized_feeds else 0)

        parent_item = parent.internalPointer()

        # 处理分类节点的子节点(FEED列表)
        if parent_item.get('type') == 'category':
            category_name = parent_item['text']
            if category_name == '未分类':
                return len(self.uncategorized_feeds)
            return len(self.category_feeds.get(category_name, []))

        return 0


    def index(self, row, column, parent=QModelIndex()):
        if not self.hasIndex(row, column, parent):
            return QModelIndex()

        if not parent.isValid():  # 根节点下的项目
            # 前3个是固定项
            if row == 0:
                return self.createIndex(row, column, {
                    'id': -1,
                    'type': 'today',
                    'text': '今日文章'
                })
            elif row == 1:
                return self.createIndex(row, column, {
                    'id': -2,
                    'type': 'unread',
                    'text': '未读文章'
                })
            elif row == 2:
                return self.createIndex(row, column, {
                    'id': -3,
                    'type': 'starred',
                    'text': '收藏文章'
                })
            else:
                # 处理分类项
                categories = sorted(self.category_feeds.keys())
                if row - 3 < len(categories):
                    category_name = categories[row - 3]
                    return self.createIndex(row, column, {
                        'id': f"cat_{hash(category_name)}",
                        'type': 'category',
                        'text': category_name
                    })
                elif row - 3 == len(categories) and self.uncategorized_feeds:
                    return self.createIndex(row, column, {
                        'id': 'uncategorized',
                        'type': 'category',
                        'text': '未分类'
                    })

        parent_item = parent.internalPointer()

        # 分类节点下的FEED项
        if parent_item.get('type') == 'category':
            category_name = parent_item['text']
            feeds = self.uncategorized_feeds if category_name == '未分类' else self.category_feeds.get(category_name, [])

            if row < len(feeds):
                feed = feeds[row]
                return self.createIndex(row, column, {
                    'id': feed.id,
                    'type': 'feed',
                    'text': feed.title,
                    'parent_category': category_name
                })

        return QModelIndex()


    def parent(self, index):
        if not index.isValid():
            return QModelIndex()

        item = index.internalPointer()

        # FEED项的父节点是分类
        if item.get('type') == 'feed':
            for row in range(self.rowCount()):
                parent_index = self.index(row, 0)
                parent_item = parent_index.data(Qt.UserRole)
                if parent_item and parent_item.get('type') == 'category' \
                        and parent_item['text'] == item['parent_category']:
                    return parent_index

        return QModelIndex()


    def hasChildren(self, parent=QModelIndex()):
        if not parent.isValid():
            return True

        parent_item = parent.internalPointer()
        return parent_item and parent_item.get('type') == 'category'


    def data(self, index, role=Qt.DisplayRole):
        if not index.isValid():
            return None

        item = index.internalPointer()

        if role == Qt.DisplayRole:
            if item['type'] == 'today':
                return f"今日文章 ({self.today_count})"
            elif item['type'] == 'unread':
                return f"未读文章 ({self.unread_counts.get('total', 0)})"
            elif item['type'] == 'starred':
                return f"收藏文章 ({self.starred_count})"
            elif item['type'] == 'feed':
                count = self.unread_counts.get(item['id'], 0)
                return f"{item['text']} ({count})" if count > 0 else item['text']
            return item['text']

        elif role == Qt.UserRole:
            return item

        elif role == Qt.DecorationRole:
            if item['type'] == 'today':
                return QIcon.fromTheme('view-calendar-day')
            elif item['type'] == 'unread':
                return QIcon.fromTheme('mail-mark-unread')
            elif item['type'] == 'starred':
                return QIcon.fromTheme('rating')
            elif item['type'] == 'category':
                return QIcon.fromTheme('folder')
            elif item['type'] == 'feed':
                return QIcon.fromTheme('application-rss+xml')

        return None


    def refresh(self):
        self.beginResetModel()
        self._setup_model_data()
        self.endResetModel()